package su2010.puz.util;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import su2010.puz.Extractor;
import su2010.puz.FeatureVector;
import su2010.puz.ImageHelper;

public class SetCreator {

	/** K: class, V: list of files for that class. */
	private Map<String, List<File>> trainingFiles;
	private Map<String, List<File>> testFiles;
	private List<String> classes;
	             
	/**
	 * 
	 * @param imgDirPath Path do directory with images.
	 * @param percentage Which percent of each image set will be set for training.
	 * @param imagesPerClass Determines how many images will be taken for every class
	 * , if 0 all images will be taken.
	 */
	public SetCreator(String imgDirPath, int percentage, int imagesPerClass) {
		init(imgDirPath, percentage, imagesPerClass);
	}
	
	/**
	 * Crawls HDD starting under dir imgDir in search for images.
	 * Each dir represents class and contains class images.
	 * 
	 * @param imgDirPath Path do directory with images.
	 * @param percentage Which percent of each image set will be set for training.
	 * @param imagesPerClass Determines how many images will be taken for every class
	 * , if 0 all images will be taken.
	 */
	private void init(String imgDirPath, int percentage, int imagesPerClass) {
		classes = new LinkedList<String>();
		trainingFiles = new HashMap<String, List<File>>();
		testFiles = new HashMap<String, List<File>>();
		
		File dir = new File(imgDirPath);
		File[] files = dir.listFiles();
		for (int i = 0; i < files.length; i++) {
			File f = files[i];
			if (f.isDirectory()) {
				String cls = f.getName();
				classes.add(cls);
				
				// List all images inside dir f and connect them with class cls
				List<File> imgs = new LinkedList<File>();
				File[] imgFiles = f.listFiles();
				for (int j = 0; j < imgFiles.length; j++) {
					File imgF = imgFiles[j];
					imgs.add(imgF);
				}
				
				// Randomize img set
				Collections.shuffle(imgs);
				
				// Split img set to training and test part
				int size = imagesPerClass == 0 || imagesPerClass>imgs.size() ? imgs.size() : imagesPerClass;
				int border = (size*percentage)/100;
				trainingFiles.put(cls, new LinkedList<File>(imgs.subList(0, border)));
				testFiles.put(cls, new LinkedList<File>(imgs.subList(border, size)));
			}
		}
	}
	
	/**
	 * @return Training set with all classes.
	 */
	public Set<FeatureVector> getFullTrainingSet(List<Extractor> extractors) {
		System.out.println("Creating training set:");
		Set<FeatureVector> training = new HashSet<FeatureVector>();
		for (String cls : classes) {
			training.addAll(getTrainingSet(cls, extractors));
		}
		
		return training;
	}
	
	/**
	 * 
	 * @return Test set with all classes.
	 */
	public Set<FeatureVector> getFullTestSet(List<Extractor> extractors) {
		System.out.println("Creating testing set:");
		Set<FeatureVector> test = new HashSet<FeatureVector>();
		for (String cls : classes) {
			test.addAll(getTestSet(cls, extractors));
		}
		
		return test;
	}
	
	public Set<FeatureVector> getTrainingSet(String cls, List<Extractor> extractors){
		return getSet(cls, trainingFiles.get(cls), extractors);
	}
	
	public Set<FeatureVector> getTestSet(String cls, List<Extractor> extractors){
		return getSet(cls, testFiles.get(cls), extractors);
	}
	
	public List<String> getClasses() {
		return classes;
	}

	private Set<FeatureVector> getSet(String cls, List<File> files, List<Extractor> extractors){
		Set<FeatureVector> set = new HashSet<FeatureVector>();
		System.out.println("Class: "+cls);
		int j = 1;
		boolean updateCache = false;
		
		try {
			Cache.instance().load("puz.cache");
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
		}
		
		for(File f : files){
			System.out.println((j++)+"/"+files.size()+" "+f.getName());
			ImageCls image = ImageHelper.readImage(f.getAbsolutePath());
			if(image!=null){
				FeatureVector fv = new FeatureVector(cls);
			
				for (Extractor extractor : extractors) {
					float[] features = Cache.instance().get(f.getAbsolutePath(), extractor.getName());
					if (features == null) {
						extractor.setImage(image);
						features = extractor.extract();
						Cache.instance().add(f.getAbsolutePath(), extractor.getName(), features);
						updateCache = true;
					}
					fv.appendFeatures(features);
				}
				image.flush();
				Runtime.getRuntime().freeMemory();
				set.add(fv);
			}
		}
		
		if (updateCache) {
			try {
				Cache.instance().save("puz.cache");
			} catch (FileNotFoundException e) {
				e.printStackTrace();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
		
		return set;
	}
}
